# -*- coding: utf-8 -*-

from fastapi import FastAPI, Request, HTTPException
from transformers import AutoTokenizer, AutoModelForCausalLM
import uvicorn
import json
import datetime
import torch
import pynvml  # 用于检查 GPU 状态

# 设置设备参数
DEVICE_IDS = ["0", "1", "2", "3", "4", "5","6"]  # 可用的 GPU 列表

# 初始化 pynvml
pynvml.nvmlInit()

# 清理 GPU 内存函数
def torch_gc():
    if torch.cuda.is_available():  # 检查是否可用 CUDA
        for device_id in DEVICE_IDS:
            with torch.cuda.device(f"cuda:{device_id}"):  # 指定 CUDA 设备
                torch.cuda.empty_cache()  # 清空 CUDA 缓存
                torch.cuda.ipc_collect()  # 收集 CUDA 内存碎片

# 检查 GPU 状态
def check_gpu_status(device_id):
    """
    检查指定 GPU 的状态
    :param device_id: GPU 设备 ID（字符串，如 "0"）
    :return: 是否可用（布尔值）
    """
    handle = pynvml.nvmlDeviceGetHandleByIndex(int(device_id))
    mem_info = pynvml.nvmlDeviceGetMemoryInfo(handle)
    util_info = pynvml.nvmlDeviceGetUtilizationRates(handle)

    # 显存使用率超过 90% 或 GPU 利用率超过 90% 则认为 GPU 不可用
    if mem_info.used / mem_info.total > 0.9 or util_info.gpu > 90:
        return False
    return True

# 获取可用 GPU
def get_available_gpu():
    """
    获取一个可用的 GPU
    :return: 可用的 GPU 设备 ID（字符串，如 "0"）
    """
    for device_id in DEVICE_IDS:
        if check_gpu_status(device_id):
            return device_id
    return None  # 如果没有可用 GPU，返回 None

# 创建 FastAPI 应用
app = FastAPI()

# 处理 POST 请求的端点
@app.post("/")
async def create_item(request: Request):
    global model, tokenizer  # 声明全局变量以便在函数内部使用模型和分词器
    json_post_raw = await request.json()  # 获取 POST 请求的 JSON 数据
    json_post = json.dumps(json_post_raw)  # 将 JSON 数据转换为字符串
    json_post_list = json.loads(json_post)  # 将字符串转换为 Python 对象
    prompt = json_post_list.get('prompt')  # 获取请求中的提示

    # 获取可用 GPU
    device_id = get_available_gpu()
    if device_id is None:
        raise HTTPException(status_code=503, detail="所有 GPU 都不可用，请稍后重试")

    # 将模型和数据移动到可用 GPU
    device = f"cuda:{device_id}"
    model.to(device)

    messages = [
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user", "content": prompt}
    ]

    # 调用模型进行对话生成
    input_ids = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    model_inputs = tokenizer([input_ids], return_tensors="pt").to(device)
    generated_ids = model.generate(model_inputs.input_ids, max_new_tokens=512)
    generated_ids = [
        output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
    ]
    response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
    now = datetime.datetime.now()  # 获取当前时间
    time = now.strftime("%Y-%m-%d %H:%M:%S")  # 格式化时间为字符串
    # 构建响应 JSON
    answer = {
        "response": response,
        "status": 200,
        "time": time,
        "device_id": device_id  # 返回使用的 GPU 设备 ID
    }
    # 构建日志信息
    log = "[" + time + "] " + '", prompt:"' + prompt + '", response:"' + repr(response) + '", device_id:"' + device_id + '"'
    print(log)  # 打印日志
    torch_gc()  # 执行 GPU 内存清理
    return answer  # 返回响应

# 主函数入口
if __name__ == '__main__':
    # 加载预训练的分词器和模型
    model_name_or_path = '/data/chenbo/wzh/Qwen/Qwen2.5-14B-Instruct-GPTQ-Int4'
    tokenizer = AutoTokenizer.from_pretrained(
        model_name_or_path,
        use_fast=False,
        local_files_only=True  # 强制加载本地文件
    )


    model = AutoModelForCausalLM.from_pretrained(
        model_name_or_path,
        device_map="auto",  # 自动分配模型层到 GPU
        torch_dtype=torch.bfloat16
    )

    # 启动 FastAPI 应用
    # 用 6006 端口可以将 autodl 的端口映射到本地，从而在本地使用 api
    uvicorn.run(app, host='0.0.0.0', port=6006, workers=1)  # 启动单个工作进程